import numpy as np
import random
from sklearn.model_selection import KFold
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error #MSE
from sklearn.metrics import mean_absolute_error #MAE
from sklearn.metrics import r2_score #R2

from xgboost import XGBRegressor as XGBR
import pandas as pd
import fig

DATA_PATH = "data/xgboost-dmax/"
elements = dict()
# 读取数据
def load_data(path):
    data = pd.read_csv(path, delimiter=r",")
    values = data.values[:, 0:1].ravel()
    valid = False
    X = []
    for data_idx in range(len(values)):
        alloy, alpha_st, num_st = [], -1, -1
        value = values[data_idx]
        leng = len(value)
        for i in range(leng):
            if value[i].isalpha() :
                if alpha_st == -1:
                    alpha_st = i
                if num_st != -1 :
                    a = value[num_st:i]
                    alloy.append(float(a))
                    num_st = -1
                if i == leng-1:
                    alloy.append(value[num_st:leng])
            else:
                if num_st == -1:
                    num_st = i
                if alpha_st != -1:
                    a = value[alpha_st:i]
                    alloy.append(a)
                    alpha_st = -1
        alloy.append(float(value[num_st:leng]))

        node = np.zeros((100, 1))
        used_idx = []

        for i in range(0, len(alloy), 2):
            idx = elements.get(alloy[i], len(elements))
            if idx == len(elements):
                elements[alloy[i]] = idx
            used_idx.append(idx)
            node[idx] = alloy[i+1]
        X.append(node.ravel())
    return np.array(X), np.array(data.values[:, -1].ravel())

def generate_dataset():
    path = DATA_PATH + "all-dmax2.csv"
    X, Y = load_data(path)

    elen = len(elements)
    scaler = MinMaxScaler(feature_range=(0, 1))
    X = scaler.fit_transform(X)[:, 0:elen]
    for i in range(len(Y)):
        Y[i] = np.log(Y[i])
    #x_train = np.array(x_train) / 100.0
    #x_valid = np.array(x_valid) / 100.0
    return X, Y

def split_index(n_total, ratio, shuffle=True):
    idx_list = list(range(0, n_total))
    offset = int(n_total * ratio)
    random.shuffle(idx_list)
    sublist1 = idx_list[:offset]
    sublist2 = idx_list[offset:]
    return sublist1, sublist2

def get_alloy_systems(X):
    alloy_systems = []
    for a in X:
        idx = [i for i, e in enumerate(a) if e != 0]
        if idx not in alloy_systems:
            alloy_systems.append(idx)
    return alloy_systems

def split_by_alloy_system_index(X, Y, alloy_system_idx):
    v_alloy_idx = []
    for i in range(len(X)):
        alloy_item = X[i]
        for alloy_sys_item in alloy_system_idx:
            if np.count_nonzero(alloy_item) == len(alloy_sys_item):
                tag = True
                for e in alloy_sys_item:
                    if alloy_item[e] == 0:
                        tag = False
                        break
                if tag:
                    v_alloy_idx.append(i)
    return v_alloy_idx

#输入划分多少，以及是否按照alloy system划分，得到test index
def split_test_data(X, Y, split_rate):
    alloy_systems = get_alloy_systems(X)
    v_sys_idx, t_sys_idx = split_index(len(alloy_systems), split_rate, shuffle=True)
    valid_alloy_system_idx = []
    for idx in v_sys_idx:
        valid_alloy_system_idx.append(alloy_systems[idx])
    test_alloy_sys_index = split_by_alloy_system_index(X, Y, valid_alloy_system_idx)
    out_data = pd.DataFrame(data=test_alloy_sys_index, columns=["idx"])
    out_data.to_csv(DATA_PATH + "test_alloy_sys_idx.csv", index=False)

    test_alloy_index, train_alloy_index = split_index(len(Y), split_rate, shuffle=True)
    out_data = pd.DataFrame(data=test_alloy_index, columns=["idx"])
    out_data.to_csv(DATA_PATH + "test_alloy_idx.csv", index=False)

def kfold_alloy_system_valid(X, Y, model):
    alloy_systems = get_alloy_systems(X)
    n_splits = 10
    kf = KFold(n_splits, shuffle=True)
    R2, MAE = 0, 0
    for train_index, valid_index in kf.split(alloy_systems):
        valid_alloy_system_idx = []
        for idx in valid_index:
            valid_alloy_system_idx.append(alloy_systems[idx])

        v_alloy_sys_index = split_by_alloy_system_index(X, Y, valid_alloy_system_idx)
        x_valid, y_valid = np.array(X[v_alloy_sys_index]), np.array(Y[v_alloy_sys_index])

        x_train, y_train = [], []
        for i in range(len(Y)):
            if i not in v_alloy_sys_index:
                x_train.append(X[i])
                y_train.append(Y[i])

        model.fit(x_train, y_train)  # 训练
        y_pred = model.predict(x_valid)
        R2 = R2 + r2_score(y_valid, y_pred)
        MAE = MAE + mean_absolute_error(y_valid, y_pred)
    print(f'kfold alloy systems r2 {R2 / n_splits} mae {MAE / n_splits}')


def kfold_alloy_valid(X, Y, model):
    n_splits = 10
    kf = KFold(n_splits, shuffle=True)
    R2, MAE = 0, 0
    for train_index, valid_index in kf.split(X):
        x_train, y_train = X[train_index], Y[train_index]
        x_valid, y_valid = X[valid_index], Y[valid_index]

        model.fit(x_train, y_train)  # 训练
        y_pred = model.predict(x_valid)

        R2 = R2 + r2_score(y_valid, y_pred)
        MAE = MAE + mean_absolute_error(y_valid, y_pred)
    print(f'kfold r2 {R2 / n_splits} mae {MAE / n_splits}')


xgbr = XGBR(n_estimators=300,
            max_depth=8,
            min_child_weight=1,
            learning_rate=0.1,
            subsample=0.8,
            colsample_bytree=0.8,
            gamma=0
            )

def set_elements():
    data = pd.read_csv("data/elements.csv", delimiter=r",")
    A = data.values[:, 0:1].ravel()
    B = data.values[:, -1].ravel()
    for i in range(len(A)):
        elements[A[i]] = B[i]

def predict(x_test):
    set_elements()
    X, Y = generate_dataset()
    xgbr.fit(X, Y)

    y_pred = xgbr.predict(x_test)
    print(len(y_pred))
    if len(y_pred) < 400:
        print(y_pred)

    y_pred = np.exp(y_pred)
    less_than1 = 0
    for yy in y_pred:
        if yy < 1:
            less_than1 = less_than1 + 1

    print(f'total {len(y_pred)} less than 1 {less_than1} {1-less_than1/len(y_pred)}')
    return y_pred

#分离出来两种test文件，输出test index
NEW_TEST_DATA = False
TEST = True
def main():
    set_elements()
    X, Y = generate_dataset()
    if NEW_TEST_DATA:
        split_test_data(X, Y, 0.1)

    #data loading
    test_aly_idx = pd.read_csv(DATA_PATH + "test_alloy_idx.csv", delimiter=r",").values
    test_aly_sys_idx = pd.read_csv(DATA_PATH + "test_alloy_sys_idx.csv", delimiter=r",").values

    w = len(elements)
    X_alloy_test, Y_alloy_test = X[test_aly_idx].reshape(len(test_aly_idx),w), Y[test_aly_idx]
    X_alloy_sys_test = np.array(X[test_aly_sys_idx]).reshape(len(test_aly_sys_idx), w)
    Y_alloy_sys_test = np.array(Y[test_aly_sys_idx])
    X_alloy, Y_alloy, X_alloy_sys, Y_alloy_sys = [], [], [], []
    for i in range(len(Y)):
        if i not in test_aly_idx:
            X_alloy.append(X[i])
            Y_alloy.append(Y[i])
        if i not in test_aly_sys_idx:
            X_alloy_sys.append(X[i])
            Y_alloy_sys.append(Y[i])

    X_alloy, Y_alloy = np.array(X_alloy), np.array(Y_alloy)
    X_alloy_sys, Y_alloy_sys = np.array(X_alloy_sys), np.array(Y_alloy_sys)
    print(f"{len(Y_alloy)} {len(Y_alloy_sys)}")

    #valid alloy
    kfold_alloy_valid(X_alloy, Y_alloy, xgbr)
    # valid alloy system
    kfold_alloy_system_valid(X_alloy_sys, Y_alloy_sys, xgbr)

    if TEST:
        # test alloy
        xgbr.fit(X_alloy, Y_alloy)
        y_pred_alloy = xgbr.predict(X_alloy_test)
        r2 = r2_score(Y_alloy_test, y_pred_alloy)
        mae = mean_absolute_error(Y_alloy_test, y_pred_alloy)
        fig.scatter2d_1(Y_alloy_test, y_pred_alloy)
        # for i in range(len(y_pred_alloy)):
        #     print(f'pred {y_pred_alloy[i]} real {y_pred_alloy[i]}')
        print(f'data test r2 {r2} mae {mae}')

        # test alloy system
        xgbr.fit(X_alloy_sys, Y_alloy_sys)
        y_pred_aly_sys = xgbr.predict(X_alloy_sys_test)
        r2 = r2_score(Y_alloy_sys_test, y_pred_aly_sys)
        mae = mean_absolute_error(Y_alloy_sys_test, y_pred_aly_sys)
        fig.scatter2d_1(Y_alloy_sys_test, y_pred_aly_sys)
        print(f'alloy system test r2 {r2} mae {mae}')

if __name__ == '__main__':
    main()


